#!/usr/bin/env python
# encoding=utf-8

import re
import sys
from itertools import chain

from elasticsearch import Elasticsearch
from pymongo import MongoClient

from config import *

reload(sys)
sys.setdefaultencoding('utf-8')

__author__ = 'xlzd'

MONGO = MongoClient(MONGO_HOST, MONGO_PORT)
ES = Elasticsearch(ES_URL)
SPACE_ENTER = re.compile(ur'[ |\n]+')


def render_anno_db(item):
    content = item.get('content', u'')[:1000]
    content = SPACE_ENTER.sub(' ', content)
    if content:
        content = content[:110] + u'...'
    else:
        content = u'当前摘要不可用'
    return {
        'title': item['announcementTitle'],
        'description': content,
        'stockCode': item['stockCode'],
        'companyShortName': item['companyShortName'],
        'releaseTime': u'{}年{}月{}日'.format(*item['announcementTime'].split(' ')[0].split('-')),
        'id': item['_id'],
        'source': item.get('source', ''),
        'type': item['announcementType'],
        'file': item['file']
    }


def render_anno_es(item, keys):
    if not item.get('_source', ''):
        return None

    if not keys:
        content = item['_source']['content'][:1000]
        content = SPACE_ENTER.sub(' ', content)
        content = content[:110] + u'...'
    else:
        content = SPACE_ENTER.sub('', item['_source']['content'])
        start = content.find(keys[-1])

        real_start = max(0, start - 10)
        real_end = real_start + len(keys[-1]) + 100
        content = content[real_start:real_end]
        for key in keys:
            content = content.replace(key, u'<span class="key">%s</span>' % key)

    result = {
        'title': item['_source']['title'],
        'description': content,
        'stockCode': item['_source']['stockCode'],
        'companyShortName': item['_source']['companyShortName'],
        'releaseTime': u'{}年{}月{}日'.format(*item['_source']['time'].split(' ')[0].split('-')),
        'id': item['_source']['id'],
        'source': item['_source']['source'],
        'type': item['_source']['type'],
    }
    return result


def select_db(filter={}, page=1, render=True, db='crawler_company_all', coll='stockAnnouncement'):
    coll = MONGO[db][coll]
    data = coll.find(filter).sort('updateTime', -1).skip(10 * page - 10).limit(10)
    if render:
        return [render_anno_db(x) for x in data]
    return [x for x in data]


def select_collect(user, page=1):
    coll = MONGO['lawyer_asst']['collect']
    data = coll.find({'userUuid': user}).sort('collectTime', -1)
    total = data.count()
    data = data.skip(10 * page - 10).limit(10)
    return total, [x for x in data]


def select_db_single(coll, filter, db='crawler_company_all'):
    return MONGO[db][coll].find_one(filter)


def remove_from_db(db, coll, filter):
    return MONGO[db][coll].find_one_and_delete(filter)


def select_es(keys, page=1, count=10):
    if isinstance(keys, basestring):
        keys = [keys]
    query = [{"match": {"content": {"query": key, "type": "phrase"}}} for key in keys]
    result = ES.search(
        index='stock_announcement',
        body={
            'from': count * page - count,
            'size': count,
            'query': {"bool": {"must": query}}
        }
    )
    total = result['hits']['total']
    return total, [render_anno_es(item, keys) for item in result['hits']['hits']]


def select_es_high(param, page=1, count=10):
    render = lambda k, v: {"match": {k: {"query": v, "type": "phrase"}}}

    must = [render('title', v) for v in param['t-a']] + [render('content', v) for v in param['c-a']]
    should = [render('title', v) for v in param['t-o']] + [render('content', v) for v in param['c-o']]
    must_not = [render('title', v) for v in param['t-n']] + [render('content', v) for v in param['c-n']]

    search_body = {
        'from': count * page - count,
        'size': count,
        'query': {
            "bool": {
                "must": must,
                "should": should,
                "must_not": must_not,
            }
        }
    }
    # import json
    # print json.dumps(search_body, ensure_ascii=False, indent=4)
    result = ES.search(index='stock_announcement', body=search_body)
    total = result['hits']['total']

    keys = param['c-o'] + param['c-a']
    return total, [render_anno_es(item, keys=keys) for item in result['hits']['hits']]


def insert_into_mongo(db, coll, data):
    MONGO[db][coll].insert_one(data)


if __name__ == '__main__':
    import json

    print json.dumps(select_es_high({
        "c-o": [],
        "c-n": [u"厦门"],
        "c-a": [u"分公司"],
        "t-a": [u"公司", u"北京", u"公告"],
        "t-n": [u"报告"],
        "t-o": []
    }), ensure_ascii=False, indent=4, sort_keys=True)
    # print json.dumps(select_es(u'阳光城集团股份有限公司'), ensure_ascii=False, indent=4, sort_keys=True)
